// ProjectInjection.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"

#include <string>

#include <Windows.h>
#include <Psapi.h>
#include <vector>

#include <boost\scoped_array.hpp>

#include "../HookProc/WinImplExpection.h"


#include "../HookProc/asm86.h"


using namespace sd::tools::asmcommand::x86;

/*
		0119B167 68 00 FF 00 FF       push        0FF00FF00h  
		0119B16C E8 73 81 00 00       call        __imp__LoadLibraryW@4 (11A32E4h)  
		0119B171 50                   push        eax  
		0119B172 68 17 FF 00 FF       push        0FF00FF17h  
		0119B177 E8 6C 81 00 00       call        __imp__GetProcAddress@8 (11A32E8h)  
		0119B17C FF D0                call        eax  
*/

struct Injection
{
	Injection( void* placment ): pushVal1_( (DWORD)placment + offsetof( Injection, Injection::buf_ ) ), callNear1_( 0 ), pushVal2_( 1 ), callNear2_(0)
	{
		//TODO Add chech for sityation when Kernel32 on not standard address
		HMODULE lib = LoadLibrary( _T("Kernel32.dll") );
		FARPROC proc1 = GetProcAddress( lib, "LoadLibraryW" );
		FARPROC proc2 = GetProcAddress( lib, "GetProcAddress" );
		callNear1_.val_ = CalculateRelativAddres( (char*)placment, proc1 ) - offsetof( Injection, Injection::callNear1_ ) - sizeof(Injection::callNear1_);
		callNear2_.val_ = CalculateRelativAddres( (char*)placment, proc2 ) - offsetof( Injection, Injection::callNear2_ ) - sizeof(Injection::callNear2_);;
		buf_[0] = 0;
	}
	void setLib( const std::wstring& libName )
	{
		wcscpy_s( buf_, libName.c_str() );
	}
private:
	pushVal pushVal1_;
	callNear callNear1_;
	pushVal pushVal2_;
	pushEax pushEax_;
	callNear callNear2_;
	callEax callEax_;
	ret ret_;
	int3 int3_;
	wchar_t buf_[1024];
};


class ProcessInjectorBase
{
public:
	ProcessInjectorBase(): injection_( 0 ), process_( 0 ), interAddr_( 0 ){}

	void attach( const std::wstring& libName )
	{
		if ( !injection_ )
			throw std::logic_error( "Internal Error in class ProcessInjectorBase: injection_ is free" );

		injection_->setLib( libName );
		DWORD real=0;

		BOOL result = WriteProcessMemory( process_, interAddr_, injection_, sizeof(Injection), &real);

		if ( !result )
			throw sd::tools::win::WinImplException( _T("Error on call WriteProcessMemory" ) );
		
		DWORD lpThread;

		HANDLE thread = CreateRemoteThread( process_, 0, 1024*1024, (LPTHREAD_START_ROUTINE)interAddr_,0,CREATE_SUSPENDED,&lpThread);

		
	}

	virtual ~ProcessInjectorBase()=0
	{
		deattachFromProcess();
	}

protected:
	void attachToProcess( HANDLE process )
	{
		process_ = process;

		interAddr_ = VirtualAllocEx( process, 0, sizeof(Injection), MEM_COMMIT, PAGE_EXECUTE_READWRITE );
		if( !interAddr_ )
			throw sd::tools::win::WinImplException( _T("Error on call VirtualAllocEx" ) );

		injection_ = new Injection( interAddr_ );
	}

	void deattachFromProcess()
	{
		if( process_ && interAddr_ )
			VirtualFreeEx( process_, interAddr_, sizeof(Injection), MEM_RELEASE );
		if( injection_ )
			delete injection_;
	}


private:
	HANDLE     process_;
	Injection* injection_;
	void* interAddr_;
};


class ProcessInjectorCreate : public ProcessInjectorBase
{
public:
	
	~ProcessInjectorCreate(){ resumeThread(); }
	ProcessInjectorCreate( const std::wstring& cmdLine ): locked_(false)
	{
		
		STARTUPINFO startupInfo;

		ZeroMemory(&startupInfo,sizeof(STARTUPINFO));
		startupInfo.cb=sizeof(STARTUPINFO);

		boost::scoped_array< wchar_t > buf1( new wchar_t[cmdLine.size()+1] );
		wcscpy_s( buf1.get(), cmdLine.size()+1, cmdLine.c_str()  );

		BOOL result = CreateProcessW( NULL, buf1.get(), 0, 0, FALSE, CREATE_SUSPENDED, 0, 0, &startupInfo, &procInfo );

		if( !result )
			throw sd::tools::win::WinImplException( std::wstring( _T("Error on start process \"") ) + cmdLine + _T( "\"" ) );

		locked_ = true;

		attachToProcess( procInfo.hProcess );
	}

	void resumeThread()
	{
		if( locked_ && procInfo.hThread )
		{
			ResumeThread( procInfo.hThread );
			locked_ = false;
		}
	}
private:
	bool locked_;
	PROCESS_INFORMATION procInfo;


};




class ProcessInjectorAttach : public ProcessInjectorBase
{
public:
	ProcessInjectorAttach( DWORD process ){}
};





void printHelp()
{
	std::cout << "Load library in onother process address space. Usage:\n program [libpath]libname -s cmd_command\n program [libpath]libname process_handle"<<std::endl;
}



int _tmain(int argc, _TCHAR* argv[])
{
	if( argc < 3 )
	{
		printHelp();
		return -1;
	}

	std::wstring libName = argv[1];

	std::wstring param = argv[2];

	try
	{
		std::auto_ptr< ProcessInjectorBase > process;

		if( param==L"-s" || param==L"-S" )
		{
			if( argc < 4 )
			{
				printHelp();
				return -1;
			}

			std::wstring params = argv[3];

			for( int i = 4; i < argc; ++i )
			{
				params+=L" ";
				params+= argv[4];

			}
			process.reset( new ProcessInjectorCreate( params ) );
		}else
		{
		
			int handle = _wtoi( argv[3] );

			if( handle <= 0 )
			{
				printHelp();
				return -1;
			}

			process.reset( new ProcessInjectorAttach( handle ) );

		}

		process->attach( libName );
	}catch( std::exception& exp )
	{
		std::cout << exp.what();
		return -2;
	}


	return 0;
}


